package com.yangxk.kadmin.web.common.datascope;

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.dialect.mysql.parser.MySqlStatementParser;
import com.yangxk.kadmin.common.util.ValidUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.stereotype.Component;

import java.lang.reflect.Field;
import java.util.List;
import java.util.Properties;

@Intercepts({
    @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
    @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class})
})
@Component
@Slf4j
public class DataScopeInterceptor implements Interceptor {
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        if (DataScopeHolder.get() == null) {
            return invocation.proceed();
        }
        Object[] args = invocation.getArgs();

        MappedStatement ms = (MappedStatement) args[0];
        Object parameter = args[1];
        BoundSql boundSql = null;
        if (args.length == 4) {
            boundSql = ms.getBoundSql(parameter);
        } else {
            boundSql = (BoundSql) args[5];
        }

        // 原始sql
        String sql = boundSql.getSql();
        DataScopeParam dataScopeParam = DataScopeHolder.get();
        // 改写后的SQL
        String newSql = getSql(dataScopeParam, sql);

        // 反射修改boundSql中的SQL
        Field field = boundSql.getClass().getDeclaredField("sql");
        field.setAccessible(true);
        field.set(boundSql, newSql);

        return invocation.proceed();
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {

    }

    private String getData(List param) {
        StringBuilder sb = new StringBuilder();
        for (Object o : param) {
            sb.append(o);
            sb.append(",");
        }

        String res = sb.toString();
        return res.substring(0, res.length() - 1);
    }

    private String getSql(DataScopeParam dataScopeParam, String sql) {
        log.info("数据权限拦截器, sql: {}", sql);
        // 利用druid 的sql parser 添加条件
        MySqlStatementParser parser = new MySqlStatementParser(sql);
        SQLSelectStatement select = (SQLSelectStatement)parser.parseStatement();
        SQLSelectQueryBlock queryBlock = select.getSelect().getQueryBlock();

        // 需要增加的条件
        String addWhere = "";

        // 单表查询
        if (queryBlock.getFrom() instanceof SQLExprTableSource) {
            String alias = getTableAlias(dataScopeParam, queryBlock.getFrom());
            addWhere = getWhere(alias, dataScopeParam);
            select.addWhere(SQLUtils.toSQLExpr(addWhere));
        }

        // 连接查询(内连接 外连接)
        else if (queryBlock.getFrom() instanceof SQLJoinTableSource) {
            String alias = getTableAlias(dataScopeParam, queryBlock.getFrom());

            addWhere = getWhere(alias, dataScopeParam);

            select.addWhere(SQLUtils.toSQLExpr(addWhere));
        }

        // 被分页mapper改写后查询总数的查询是一个子查询, 需要提取出子查询进行判断
        else if (queryBlock.getFrom() instanceof SQLSubqueryTableSource) {
            SQLSubqueryTableSource from = (SQLSubqueryTableSource) queryBlock.getFrom();
            if (!from.getAlias().equals("tmp_count"))
                throw new RuntimeException("数据权限只支持单表查询和连接查询");

            queryBlock = from.getSelect().getQueryBlock();
            if (!(queryBlock.getFrom() instanceof SQLJoinTableSource))
                throw new RuntimeException("数据权限只支持单表查询和连接查询");

            String alias = getTableAlias(dataScopeParam, queryBlock.getFrom());

            addWhere = getWhere(alias, dataScopeParam);

            queryBlock.addCondition(SQLUtils.toSQLExpr(addWhere));
        }

        // 其他类型的查询不支持
        else {
            throw new RuntimeException("数据权限只支持单表查询和连接查询");
        }

        return select.toString();
    }

    private String getWhere(String alias, DataScopeParam dataScopeParam) {
        String addWhere = "";
        if (!ValidUtil.isEmpty(alias))
            addWhere = alias + ".";
        addWhere += dataScopeParam.getFieldName() + " in (" + getData(dataScopeParam.getLimitNumbers()) + ")";

        return addWhere;
    }

    private String getTableAlias(DataScopeParam dataScopeParam, SQLTableSource ss) {
        // DataScope注解中设置的表别名(优先级最高)
        String paramAlias = dataScopeParam.getTableAlias();

        String fromAlias = "";
        if (ss instanceof SQLExprTableSource) {
            fromAlias = ss.getAlias();
        }
        else if (ss instanceof SQLJoinTableSource) {
            fromAlias = internalAlias((SQLJoinTableSource)ss);
        }

        String alias = !ValidUtil.isEmpty(paramAlias) ? paramAlias : fromAlias;

        return alias;
    }

    private String internalAlias(SQLJoinTableSource s) {
        if (s.getLeft() instanceof SQLJoinTableSource) {
            return internalAlias((SQLJoinTableSource)s.getLeft());
        } else if (s.getLeft() instanceof SQLExprTableSource) {
            return s.getLeft().getAlias();
        }

        return "t1";
    }
}
